import os
os.environ['DJANGO_SETTINGS_MODULE']='db.settings'

from PyQt4.QtCore import *
from PyQt4.QtGui import *

from django.db.models.signals import post_save
from django.dispatch import receiver
from stock.models import * 

_Category = 0
_Pk       = 1

class TreeItem(object):
    def __init__(self, data, parent=None):
        self.childItems=[]
        self.itemData=data
        self.parentItem=parent

    def child(self, number): return self.childItems[number]

    def childCount(self): return len(self.childItems)

    def columnCount(self): return len(self.itemData)

    def data(self, column): return self.itemData[column]

    def insertChildren(self, position, count, columns):
        if position<0 or position>len(self.childItems): return False

        for row in range(count):
            data=[QVariant() for _ in range(columns)]
            self.childItems.insert(position, TreeItem(data, self))
        return True

    def insertColumns(self, position, columns):
        if position<0 or position>len(self.itemData): return False

        for column in range(columns): self.itemData.insert(position, QVariant())
        for child in self.childItems: child.insertColumns(position, columns)
        return True

    def parent(self): return self.parentItem

    def removeChildren(self, position, count):
         if position<0 or (position+count)>len(self.childItems): return False

         for row in range(count): del self.childItems[position]
         return True

    def removeColumns(self, position, columns):
        if position<0 or (position+columns)>len(self.itemData): return False

        for column in range(columns): del self.itemData[position]
        for child in self.childItems: child.removeColumns(position, columns)
        return True

    def childNumber(self):
        if self.parentItem: return self.parentItem.childItems.index(self)
        return 0

    def setData(self, column, value):
        if column<0 or column>=len(self.itemData): return False

        self.itemData[column]=value

        if column==_Pk:
            _object=ProductType.objects.get(pk=value)
            for i, child in enumerate(_object.get_children()):
                self.insertChildren(i, 1, 2)
                self.child(i).setData(0, child.name)
                self.child(i).setData(1, child.pk)

        return True

class TreeModel(QAbstractItemModel):
    self_=None

    itemChanged=pyqtSignal(str, int)

    def __init__(self, headers, data, parent=None):
        assert not TreeModel.self_
        TreeModel.self_=self

        QAbstractItemModel.__init__(self, parent)

        self.rootItem=TreeItem(headers)
        self.setupModelData(self.rootItem)

    @staticmethod
    def instance():
        if not TreeModel.self_: TreeModel.self_=TreeModel()
        return TreeModel.self_ 

    def rootIndex(self):
        return self.match(self.index(0,0), Qt.DisplayRole, "ROOT", hits=1, flags=Qt.MatchExactly|Qt.MatchWrap)[0]

    def data(self, index, role):
        if not index.isValid(): return QVariant()
        if role not in (Qt.DisplayRole, Qt.EditRole): return QVariant()

        return self.getItem(index).data(index.column())

    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if orientation==Qt.Horizontal and role==Qt.DisplayRole: return self.rootItem.data(section)
        return QVariant()

    def index(self, row, column, parent=QModelIndex()):
        if parent.isValid() and parent.column()!=0: return QModelIndex()

        childItem=self.getItem(parent).child(row)
        if childItem: return self.createIndex(row, column, childItem)
        return QModelIndex()

    def parent(self, index):
        if not index.isValid(): return QModelIndex()

        parentItem=self.getItem(index).parent()
        if parentItem==self.rootItem: return QModelIndex()
        return self.createIndex(parentItem.childNumber(), 0, parentItem)

    def rowCount(self, parent=QModelIndex()): return self.getItem(parent).childCount()

    def columnCount(self, parent=QModelIndex()): return self.rootItem.columnCount()

    def flags(self, index):
        if not index.isValid(): return 0
        return Qt.ItemIsEditable | Qt.ItemIsEnabled | Qt.ItemIsSelectable

    def setData(self, index, value, role=Qt.EditRole):
        if role!=Qt.EditRole: return False

        result=self.getItem(index).setData(index.column(), value)
        if result: self.dataChanged.emit(index, index)
        return result

    def setHeaderData(self, section, orientation, value, role=Qt.EditRole):
        if role!=Qt.EditRole or orientation!=Qt.Horizontal: return False

        result=self.rootItem.setData(section, value)
        if result: self.headerDataChanged.emit(orientation, section, section)
        return result

    def insertColumns(self, position, columns, parent=QModelIndex()):
        self.beginInsertColumns(parent, position, position+columns-1)
        success=self.rootItem.insertColumns(position, columns)
        self.endInsertColumns()
        return success

    def removeColumns(self, position, columns, parent=QModelIndex()):
        self.beginRemoveColumns(parent, position, position+columns-1)
        success=self.rootItem.removeColumns(position, columns)
        self.endRemoveColumns()

        if self.rootItem.columnCount()==0: self.removeRows(0, self.rowCount())

        return success

    def insertRows(self, position, rows, parent=QModelIndex()):
        self.beginInsertRows(parent, position, position+rows-1)
        success=self.getItem(parent).insertChildren(position, rows, self.rootItem.columnCount())
        self.endInsertRows()
        return success

    def removeRows(self, position, rows, parent=QModelIndex()):
        self.beginRemoveRows(parent, position, position+rows-1)
        success=self.rootItem.removeChildren(position, rows)
        self.endRemoveRows()
        return success

    def setupModelData(self, parent):
        for p in ProductType.objects.root_nodes():
            parent.insertChildren(parent.childCount(), 1, self.rootItem.columnCount())
            parent.child(parent.childCount()-1).setData(0, p.name)
            parent.child(parent.childCount()-1).setData(1, p.pk)

    def getItem(self, index):
        if index.isValid():
            item=index.internalPointer()
            if item: return item
        return self.rootItem

    @receiver(post_save, sender=ProductType)
    def djangpChanged(sender, **kwargs):
        TreeModel.instance().djangoItemChanged(sender, **kwargs)

    def djangoItemChanged(self, sender, **kwargs):
        instance=kwargs['instance']
        self.itemChanged.emit(instance.__class__.__name__, instance.pk)

        parentInstance=instance.get_ancestors(ascending=True)[0]

        matches=self.match(self.index(0,0), Qt.DisplayRole, parentInstance.name, hits=-1, flags=Qt.MatchRecursive|Qt.MatchStartsWith|Qt.MatchWrap)

        parentIndex=None
        parent=None
        for index in matches:
            parentIndex=index
            parent=index.internalPointer()
            if int(parent.data(1))==parentInstance.pk: break

        if parent==None:
            print "Failed to match!"
            return

        pks=[child.pk for child in parentInstance.get_children()]
        position=pks.index(instance.pk)

        if not self.insertRow(position, parentIndex): 
            print "Failed to insert row"
            return

        categoryIndex=self.index(position, _Category, parentIndex)
        self.setData(categoryIndex, instance.name, Qt.EditRole)

        pkIndex=self.index(position, _Pk, parentIndex)
        self.setData(pkIndex, instance.pk, Qt.EditRole)

class MainWindow(QMainWindow):
    def __init__(self, parent=None):
        QMainWindow.__init__(self, parent)

        self.setWindowTitle("Editable Tree Model")
        self.resize(640,480)
        self.statusBar()

        self.model=TreeModel(["Category", "pk"], "")

        self.view=QTreeView(self)
        self.view.setModel(self.model)
        self.view.setRootIndex(self.model.rootIndex())
        for column in range(self.model.columnCount()): self.view.resizeColumnToContents(column)
        self.view.setAlternatingRowColors(True)
        self.view.setSelectionMode(QAbstractItemView.SingleSelection)
        self.view.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.view.setContextMenuPolicy(Qt.CustomContextMenu)
        self.view.customContextMenuRequested.connect(self.showContextMenu)
        self.setCentralWidget(self.view)

        self.createActions()
        self.createMenus()

    def createActions(self):
        self.exitAction=QAction("&Exit", self, triggered=self.close)
        self.addTopLevelCategoryAction=QAction("Add Top Level Category...", self, triggered=self.addTopLevelCategory, shortcut="Ctrl+T")
        self.addCategoryAction=QAction("Add Category...", self, triggered=self.addCategory, shortcut="Ctrl+N")
        self.deleteCategoryAction=QAction("Delete Category...", self)

    def createMenus(self):
        self.fileMenu=self.menuBar().addMenu("&File")
        self.fileMenu.addAction(self.exitAction)

        self.actionsMenu=self.menuBar().addMenu("&Actions")
        self.actionsMenu.addAction(self.addTopLevelCategoryAction)
        self.actionsMenu.addAction(self.addCategoryAction)

    def addTopLevelCategory(self):
        new, ok=QInputDialog.getText(self, "Add Category...", "Name:")
        if not ok: return

        ProductType.objects.create(name=new, parent=ProductType.objects.get(pk=1))

    def addCategory(self):
        try: index=self.view.selectedIndexes()[0]
        except IndexError: return

        new, ok=QInputDialog.getText(self, "Add Category...", "Name:")
        if not ok: return

        __index=self.model.sibling(index.row(), _Pk, index)
        parent=ProductType.objects.get(pk=int(self.model.data(__index, Qt.DisplayRole)))
        ProductType.objects.create(name=new, parent=parent)

    def showContextMenu(self, point):        
        _menu=QMenu(self)
        _menu.addAction(self.addCategoryAction)
        _menu.addAction(self.deleteCategoryAction)
        _action=_menu.exec_(self.mapToGlobal(point))

if __name__=="__main__":
    assert len(ProductType.objects.root_nodes())==1
    from sys import argv, exit
    a=QApplication(argv)
    m=MainWindow()
    m.show()
    exit(a.exec_())
    
